package com.dream.auth.security.handler;

import java.io.IOException;
import java.time.Instant;
import java.util.LinkedHashMap;
import java.util.Map;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.web.authentication.AuthenticationFailureHandler;

import com.fasterxml.jackson.databind.ObjectMapper;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class CustomTokenEndpointErrorHandler implements AuthenticationFailureHandler {

    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void onAuthenticationFailure(HttpServletRequest request, HttpServletResponse response,
            AuthenticationException exception) throws IOException, ServletException {
        
        log.error("Token endpoint error", exception);

        Map<String, Object> errorResponse = new LinkedHashMap<>();
        HttpStatus status = HttpStatus.BAD_REQUEST; // 默认状态码

        if (exception instanceof OAuth2AuthenticationException) {
            OAuth2Error error = ((OAuth2AuthenticationException) exception).getError();
            errorResponse.put(OAuth2ParameterNames.ERROR, error.getErrorCode());
            
            if (error.getDescription() != null) {
                errorResponse.put(OAuth2ParameterNames.ERROR_DESCRIPTION, error.getDescription());
            }
            if (error.getUri() != null) {
                errorResponse.put(OAuth2ParameterNames.ERROR_URI, error.getUri());
            }

            // 根据错误类型设置不同的状态码
            switch (error.getErrorCode()) {
                case "invalid_client":
                    status = HttpStatus.UNAUTHORIZED;
                    break;
                case "invalid_grant":
                    status = HttpStatus.UNAUTHORIZED;
                    break;
                case "invalid_scope":
                    status = HttpStatus.BAD_REQUEST;
                    break;
                case "unauthorized_client":
                    status = HttpStatus.FORBIDDEN;
                    break;
                default:
                    status = HttpStatus.BAD_REQUEST;
            }
        } else {
            errorResponse.put(OAuth2ParameterNames.ERROR, "server_error");
            errorResponse.put(OAuth2ParameterNames.ERROR_DESCRIPTION, "An error occurred while processing the request");
            status = HttpStatus.INTERNAL_SERVER_ERROR;
        }

        // 添加时间戳
        errorResponse.put("timestamp", Instant.now().toString());
        // 添加请求路径
        errorResponse.put("path", request.getRequestURI());

        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        response.setStatus(status.value());
        response.setCharacterEncoding("UTF-8");
        
        this.objectMapper.writeValue(response.getWriter(), errorResponse);
    }
} 